import os
from datetime import datetime
from openai import OpenAI
import streamlit as st
from langchain_community.tools import TavilySearchResults

# 配置信息集中管理
CONFIG = {
    "tavily_api_key": "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",
    "model_endpoint": "https://api-inference.modelscope.cn/v1/",
    "model_token": "b011f8fd-99a5-4633-8bf5-fef46cf7bd90",
    "model_name": "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B",
    "model_name_V3": "deepseek-ai/DeepSeek-V3"
}

client = OpenAI(
    base_url=CONFIG["model_endpoint"],
    api_key=CONFIG["model_token"],
)
os.environ["TAVILY_API_KEY"] = CONFIG["tavily_api_key"]

# 自定义CSS样式
STYLES = """
<style>
.think-block {
    background-color: #f0f0f0;
    border-left: 4px solid #4CAF50;
    padding: 1rem;
    margin: 1rem 0;
    font-style: italic;
    color: #555;
    border-radius: 0 5px 5px 0;
}
.container-border {
    border: 1px solid #e0e0e0;
    border-radius: 8px;
    padding: 1rem;
    margin-bottom: 1rem;
}
</style>
"""

def get_web_search_results(query: str, max_retry: int = 1) -> str:
    """获取网络搜索结果并进行格式化"""

    tool = TavilySearchResults(
        max_results=5,
        search_depth="advanced",
        include_answer=True,
        include_raw_content=True,
        include_images=False,
        # include_domains=[...],
        # exclude_domains=[...],
        # name="...",            # overwrite default tool name
        # description="...",     # overwrite default tool description
        # args_schema=...,       # overwrite default args_schema: BaseModel
    )
    for attempt in range(max_retry):
        try:
            result = tool.invoke({"query": query})
            print(result)
            formatted_result = ""
            for i, item in enumerate(result, start=1):
                formatted_result += (
                    f"[webpage {i} begin]\n"
                    f"URL: {item['url']}\n"
                    f"Content: {item['content'][:1000]}...\n"  # 限制内容长度
                    f"[webpage {i} end]\n"
                )
            return formatted_result
        except Exception as e:
            print(e)
            if attempt == max_retry - 1:
                return f"网络搜索暂时不可用，错误信息：{str(e)}"
            continue


def build_system_prompt(search_results: str, query: str) -> str:
    """构建系统提示模板"""
    current_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')

    return f"""
# 以下内容是基于用户发送的消息的搜索结果:{search_results}

在我给你的搜索结果中，每个结果都是[webpage X begin]...[webpage X end]格式的，X代表每篇文章的数字索引。
请在适当的情况下在句子末尾引用上下文。请按照引用编号[citation:X]的格式在答案中对应部分引用上下文。如果一句话源自多个上下文，请列出所有相关的引用编号，例如[citation:3][citation:5]，切记不要将引用集中在最后返回引用编号，而是在答案对应部分列出。
在回答时，请注意以下几点：
- 今天是{current_time}。
- 并非搜索结果的所有内容都与用户的问题密切相关，你需要结合问题，对搜索结果进行甄别、筛选。
- 对于列举类的问题（如列举所有航班信息），尽量将答案控制在10个要点以内，并告诉用户可以查看搜索来源、获得完整信息。优先提供信息完整、最相关的列举项；如非必要，不要主动告诉用户搜索结果未提供的内容。
- 对于创作类的问题（如写论文），请务必在正文的段落中引用对应的参考编号，例如[citation:3][citation:5]，不能只在文章末尾引用。你需要解读并概括用户的题目要求，选择合适的格式，充分利用搜索结果并抽取重要信息，生成符合用户要求、极具思想深度、富有创造力与专业性的答案。你的创作篇幅需要尽可能延长，对于每一个要点的论述要推测用户的意图，给出尽可能多角度的回答要点，且务必信息量大、论述详尽。
- 如果回答很长，请尽量结构化、分段落总结。如果需要分点作答，尽量控制在5个点以内，并合并相关的内容。
- 对于客观类的问答，如果问题的答案非常简短，可以适当补充一到两句相关信息，以丰富内容。
- 你需要根据用户要求和回答内容选择合适、美观的回答格式，确保可读性强。
- 你的回答应该综合多个相关网页来回答，不能重复引用一个网页。
- 除非用户要求，否则你回答的语言需要和用户提问的语言保持一致。
# 用户消息为：{query}"""


def generate_answer(query: str,temperature=0.6, need_network = True, need_r1 = True) -> str:
    """生成答案主流程"""
    # 获取网络搜索结果 构建提示词
    print("\n=== 正在搜索最新信息 ===")
    if need_network:
        search_results = get_web_search_results(query)
        user_prompt = build_system_prompt(search_results, query)
    else:
        search_results = ""
        user_prompt = query
    print(search_results)
    if need_r1:
        model = CONFIG["model_name"]
    else:
        model = CONFIG["model_name_V3"]
    try:
        response = client.chat.completions.create(
            model = model,
            messages = [
                {"role": "user", "content": user_prompt}
            ],
            stream = True,
            temperature = temperature,  # 降低随机性
            max_tokens = 8000
        )

        print("\n\n=== 生成回答 ===")
        done_reasoning = False
        for chunk in response:
            reasoning_chunk = chunk.choices[0].delta.reasoning_content
            answer_chunk = chunk.choices[0].delta.content
            if reasoning_chunk != '':
                print(reasoning_chunk, end='', flush=True)
                yield 'think', reasoning_chunk
            elif answer_chunk != '':
                if not done_reasoning:
                    print('\n\n === Final Answer ===\n')
                    done_reasoning = True
                print(answer_chunk, end='', flush=True)
                yield 'content', answer_chunk

    except Exception as e:
        print(e)
        print(f"\n生成回答时发生错误：{str(e)}")


def render_chat_interface():
    """渲染聊天界面"""
    st.markdown(STYLES, unsafe_allow_html=True)
    st.title("🤖 DeepSeek AI 联网助手")

    # 侧边栏配置
    with st.sidebar:
        st.header("对话设置")
        temperature = st.slider("创造力", 0.0, 1.0, 0.6, step=0.1)
        need_network = st.checkbox("是否需要联网", value=True)
        need_r1 = st.checkbox("R1深度思考", value=True)

    # 初始化会话状态
    if 'messages' not in st.session_state:
        st.session_state.messages = []

    # 显示历史消息
    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.markdown(message["content"])
    # 联网按钮

    # 用户输入
    if prompt := st.chat_input("输入您的问题"):
        # 添加用户消息

        st.session_state.messages.append({"role": "user", "content": prompt})

        # 显示用户消息
        with st.chat_message("user"):
            st.markdown(prompt)


        # 显示AI响应
        with st.chat_message("assistant"):
            response_placeholder = st.empty()
            full_response = ""
            full_think_response = ""
            think_count = 0
            # 流式显示响应
            for tag, chunk in generate_answer(prompt, temperature, need_network, need_r1):
                if 'think' == tag:
                    if think_count == 0:
                        full_think_response += f'<div class="think-block">'
                    think_count = 1
                    full_think_response += chunk
                    response_placeholder.markdown(full_think_response + "▌", unsafe_allow_html=True)
                else:
                    if think_count == 1:
                        full_think_response += '</div>'
                        response_placeholder.markdown(full_think_response, unsafe_allow_html=True)
                        think_count = 0
                    full_response += chunk
                    response_placeholder.markdown(full_think_response + full_response + "▌", unsafe_allow_html=True)



        # 保存AI消息
        st.session_state.messages.append({"role": "assistant", "content": full_response})



if __name__ == "__main__":
    render_chat_interface()


